# Title     : TODO
# Objective : TODO
# Created by: Administrator
# Created on: 2019/7/24

library(ggrepel)
library(ropls)
library(pROC)
library(egg)
library(randomForest)
library(Boruta)
library(magrittr)
library(optparse)
library(caret)
library(tidyverse)

createWhenNoExist <- function(f) {
  !dir.exists(f) && dir.create(f)
}

option_list <- list(
  make_option("--i", default = "../train.csv", type = "character", help = "metabolite data file"),
  make_option("--g", default = "../train_group.csv", type = "character", help = "sample group file")
)
opt <- parse_args(OptionParser(option_list = option_list))

options(digits = 3)
args <- commandArgs(trailingOnly = F)
scriptPath <- dirname(sub("--file=", "", args[grep("--file", args)]))
source(str_c(scriptPath, "/base.R"))

sampleInfo <- read_csv(opt$g) %>%
  rename(ClassNote = Group) %>%
  select(c("SampleID", "ClassNote"))

groups <- sampleInfo %>%
  .$ClassNote %>%
  unique()

initData <- read_csv(opt$i) %>%
  rename(Metabolite = SampleID) %>%
  gather("SampleID", "Value", -Metabolite) %>%
  spread(Metabolite, "Value")

originalColumnNames <- colnames(initData)

originalData <- initData %>%
  inner_join(sampleInfo, by = c("SampleID"))

data <- originalData %>%
  mutate(ClassNote = factor(ClassNote, levels = unique(ClassNote)))

trainData <- data %>%
  column_to_rownames("SampleID")

testSampleInfoTb <- read_csv(str_c("../test_group.csv"))

testSampleInfo <- testSampleInfoTb %>%
  rename(ClassNote = Group) %>%
  select(c("SampleID", "ClassNote")) %>%
  filter(ClassNote %in% groups)

testDataTb <- read_csv("../test.csv")
testData <- testDataTb %>%
  rename(Metabolite = SampleID) %>%
  gather("SampleID", "Value", -Metabolite) %>%
  spread(Metabolite, "Value") %>%
  select(originalColumnNames) %>%
  inner_join(testSampleInfo, by = c("SampleID")) %>%
  column_to_rownames("SampleID")

numColumns <- colnames(trainData) %>%
  keep(function(x) {
    startsWith(x, "Nubmer") || startsWith(x, "Number")
  })

testData <- testData %>%
  mutate_at(vars(numColumns), as.numeric) %>%
  mutate_at(vars(-numColumns), factor)

trainData <- trainData %>%
  mutate_at(vars(numColumns), as.numeric) %>%
  mutate_at(vars(-numColumns), factor)

common <- intersect(names(trainData), names(testData))

for (p in common) {
  if (class(trainData[[p]]) == "factor") {
    testData[[p]] <- factor(testData[[p]], levels = levels(trainData[[p]]))
  }
}

testData$ClassNote

# trainData #%>%
# select(c("HMLH1")) %>%
# str()

# testData %>%
# select(c("HMLH1")) %>%
# str()

glmRs <- glm(ClassNote ~ ., data = trainData, family = binomial(link = "logit"))

varImp <- varImp(glmRs, scale = T)
varImpDf <- varImp %>%
  rownames_to_column("Metabolite") %>%
  as_tibble() %>%
  rename(VarImp = Overall) %>%
  arrange(desc(VarImp)) %>%
  mutate(Metabolite = str_replace_all(Metabolite, "`", "")) %>%
  rename(Var = Metabolite)
write_csv(varImpDf, "LR_VarImp.csv")

Yhat <- predict(glmRs, newdata = testData, type = "response")

YDf <- Yhat %>%
  as.data.frame() %>%
  rownames_to_column("SampleID") %>%
  as_tibble() %>%
  set_colnames(c("SampleID", "Value"))

thresh = 0.5
uniq.group <- unique(originalData$ClassNote)
Yfac <- as.factor(testData$ClassNote)
YhatFac <- cut(Yhat, breaks = c(-Inf, thresh, Inf), labels = c(uniq.group[1], uniq.group[2]))
sum <- length(YhatFac)
count <- 0

for (i in 1:length(YhatFac)) {
  if (as.character(YhatFac[i]) == as.character(Yfac[i])) {
    count = count + 1
  }
}
if (count / sum < 0.5) {
  YhatFac = cut(Yhat, breaks = c(-Inf, thresh, Inf), labels = c(uniq.group[2], uniq.group[1]))
}
YhatFac <- factor(YhatFac, levels = levels(Yfac))
cTab = table(Yfac, YhatFac)
cTab
pre_summary = table(YhatFac, Yfac)

summaryTb <- pre_summary %>%
  as.data.frame() %>%
  as_tibble() %>%
  rename(Var1 = 1, Var2 = 2) %>%
  spread(Var2, "Freq")
summaryMatrix <- summaryTb %>%
  select(-"Var1") %>%
  as.matrix()
diagSum <- sum(diag(summaryMatrix))
sum <- sum(summaryMatrix)
predictive <- (diagSum / sum) %>%
  round(3)
finalSummaryTb <- summaryTb %>%
  mutate(`Model predictive accuracy` = c(predictive, "")) %>%
  rename(` ` = Var1)
write_csv(finalSummaryTb, "LR_Prediction_Summary.csv")

predictDf1 <- YDf %>%
  mutate(LR_prediction = YhatFac)

print("===")
testSampleInfo
predictDf1

predictFinalDf1 <- testSampleInfo %>%
  left_join(predictDf1, by = c("SampleID"))

predictFinalDf1

write_tsv(predictFinalDf1, "Classification_Result.txt")

predictDf <- YDf %>%
  mutate(Prediction = YhatFac)
predictFinalDf <- testSampleInfo %>%
  left_join(predictDf, by = c("SampleID")) %>%
  rename(Sample = SampleID) %>%
  rename(Probability = Value) %>%
  mutate_at(vars("Probability"), function(x) {
    ifelse(x > 0.5, x, 1 - x)
  }) %>%
  select(-c("Probability"), "Probability")
predictFinalDf
write_csv(predictFinalDf, "LR_Prediction.csv")
